# -*- coding: utf-8 -*-  
# pip install opencv-python mediapipe numpy websockets  

import cv2  
import mediapipe as mp  
import numpy as np  
import websockets  
import asyncio  
import threading  
import time  
import json  
from collections import deque  
from typing import Optional, Set  

# ============================ 配置 ============================  
TRANSPORT = "ws_server"     # WebSocket服务器广播  
HOST = "127.0.0.1"  
PORT = 8765  
SOURCE_ID = "py-pose-1"  

# 预览调试  
MIRROR_PREVIEW = False  
DEBUG_PRINT = False  
PRINT_JOINTS_EVERY_N = 0  

# 平滑（对关键点像素坐标做 EMA）  
EMA_ALPHA = 0.45  

# 下蹲（多特征融合，单位为“肩宽”）  
SQUAT_HOLD_FRAMES = 3  
SQUAT_COOLDOWN_S = 1.0  
# 髋接近膝（越接近越像下蹲）：hip_y > knee_y - delta  
SQUAT_KNEE_PROX_DELTA = 0.10      # 推荐 0.08 ~ 0.15  
# 膝角阈值（两膝都小于该角度算弯曲充分）  
SQUAT_KNEE_ANGLE_MAX = 120.0      # 推荐 100 ~ 130  
# 髋-踝垂直距离缩小（站立时通常 >1.0，深蹲显著降低）  
SQUAT_HIP_ANKLE_GAP_MAX = 0.90    # 先观察再微调  
# 辅助：髋相对肩下降量（不用作强判据）  
SQUAT_HIP_DROP_VS_SHOULDER = 0.35  

# 跳跃检测（以肩宽归一化）  
# 注意：像素 y 轴向下为正；归一化后保持此方向约定  
JUMP_VEL_UP_THRESH = -1.2       # 向上速度阈值（负值，越负越严格）  
JUMP_MIN_ASCENT = 0.40          # 抬升最小幅度（肩宽）  
JUMP_FALL_CONFIRM_VEL = 0.6     # 回落速度阈值（向下为正）  
JUMP_COOLDOWN_S = 1.0  
JUMP_SMOOTH_T = 0.10            # 速度低通滤波时间常数（秒）  

# 挥手（以肩宽归一化）  
WAVE_WINDOW_S = 1.2  
WAVE_MIN_REVERSALS = 2  
WAVE_MIN_AMPLITUDE_SHOULDER = 0.6  
WAVE_COOLDOWN_S = 1.2  

# 可见度阈值（关键点质量门槛）  
VIS_THRESH = 0.3  # 原 0.5，降低以便观察膝盖  

POINT_COLOR = (0, 200, 255)  

# ============================ WebSocket 服务器 ============================  
import asyncio  
import threading  
import json  
import time  
import websockets  

class WSServer:  
    def __init__(self, host, port, source_id="source"):  
        self.host = host  
        self.port = port  
        self.source_id = source_id  
        self.seq = 0  
        self.loop: asyncio.AbstractEventLoop | None = None  
        self.thread: threading.Thread | None = None  
        self.clients: set[websockets.WebSocketServerProtocol] = set()  
        self._ready = threading.Event()  
        self._stopping = False  

    def start(self):  
        self.thread = threading.Thread(target=self._thread_main, daemon=True)  
        self.thread.start()  
        # 等待服务器就绪，避免在 loop 未准备好时调用 send_event  
        if not self._ready.wait(timeout=5.0):  
            print("[NET] WS server startup timeout warning (will still try to continue).")  
        else:  
            print(f"[NET] WS server listening on ws://{self.host}:{self.port}")  

    def _thread_main(self):  
        asyncio.set_event_loop(asyncio.new_event_loop())  
        self.loop = asyncio.get_event_loop()  

        async def handler(ws):  # 单参版本，兼容你的 websockets 环境  
            self.clients.add(ws)  
            print(f"[NET] WS client connected. total={len(self.clients)}")  
            try:  
                async for _ in ws:  
                    pass  
            except Exception as e:  
                if DEBUG_PRINT:  
                    print(f"[NET] WS client error: {e}")  
            finally:  
                self.clients.discard(ws)  
                print(f"[NET] WS client disconnected. total={len(self.clients)}")  

        async def start_server():  
            server = await websockets.serve(handler, self.host, self.port, max_queue=None)  
            # 标记就绪  
            self._ready.set()  
            return server  

        # 启动服务器  
        server = self.loop.run_until_complete(start_server())  
        try:  
            self.loop.run_forever()  
        finally:  
            # 优雅关闭  
            server.close()  
            self.loop.run_until_complete(server.wait_closed())  
            pending = asyncio.all_tasks(loop=self.loop)  
            for task in pending:  
                task.cancel()  
            self.loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))  
            self.loop.close()  

    def _build_text(self, event_dict: dict) -> str:  
        self.seq = (self.seq + 1) & 0xFFFFFFFF  
        event_dict["seq"] = self.seq  
        event_dict["source"] = self.source_id  
        return json.dumps(event_dict, ensure_ascii=False)  

    async def _broadcast_async(self, text: str):  
        if not self.clients:  
            return  
        dead = []  
        for ws in list(self.clients):  
            try:  
                await ws.send(text)  
            except Exception:  
                dead.append(ws)  
        for ws in dead:  
            self.clients.discard(ws)  

    def send_event(self, gesture: str, details: dict | None = None):  
        # 如果事件循环未就绪，避免投递导致 "coroutine was never awaited"  
        if not self._ready.is_set() or self.loop is None or self._stopping:  
            if DEBUG_PRINT:  
                print("[NET] WS not ready; drop event:", gesture)  
            return  
        payload = {"t": time.time(), "gesture": gesture}  
        if details is not None:  
            payload["details"] = details  
        text = self._build_text(payload)  

        async def _task():  
            await self._broadcast_async(text)  

        try:  
            fut = asyncio.run_coroutine_threadsafe(_task(), self.loop)  
            # 可选：添加异常回调  
            def _cb(f: asyncio.Future):  
                exc = f.exception()  
                if exc and DEBUG_PRINT:  
                    print("[NET] WS broadcast error:", exc)  
            fut.add_done_callback(_cb)  
        except RuntimeError as e:  
            if DEBUG_PRINT:  
                print("[NET] WS loop closed; drop event:", e)  

        print(f"[GESTURE] -> {gesture}")  

    def close(self):  
        self._stopping = True  
        if self.loop and self.loop.is_running():  
            self.loop.call_soon_threadsafe(self.loop.stop)  
        if self.thread:  
            self.thread.join(timeout=2.0)  

# ============================ Pose & 函数工具 ============================  
mp_pose = mp.solutions.pose  
LMS = mp_pose.PoseLandmark  

def ema(prev, cur, alpha):  
    if prev is None:  
        return cur  
    return prev * (1 - alpha) + cur * alpha  

def compute_keypoints_xy(landmarks, w, h):  
    pts = [(l.x * w, l.y * h) for l in landmarks]  
    return np.array(pts, dtype=np.float32)  

def shoulder_center_and_width(pts):  
    l_sh = pts[LMS.LEFT_SHOULDER]  
    r_sh = pts[LMS.RIGHT_SHOULDER]  
    center = 0.5 * (l_sh + r_sh)  
    width = float(np.linalg.norm(l_sh - r_sh) + 1e-6)  
    return center, width  

def angle_deg(a, b, c):  
    """  
    计算关节角度 ∠ABC（B为顶点）  
    输入/输出：二维 np.array([x,y]), 返回度数  
    """  
    v1 = a - b  
    v2 = c - b  
    n1 = np.linalg.norm(v1) + 1e-6  
    n2 = np.linalg.norm(v2) + 1e-6  
    cosv = np.clip(np.dot(v1, v2) / (n1 * n2), -1.0, 1.0)  
    return float(np.degrees(np.arccos(cosv)))  

def update_wave(history_dq, t_now, wrist_x_norm):  
    history_dq.append((t_now, wrist_x_norm))  
    while history_dq and (t_now - history_dq[0][0]) > WAVE_WINDOW_S:  
        history_dq.popleft()  
    if len(history_dq) < 5:  
        return 0, 0.0  
    xs = np.array([x for _, x in history_dq], dtype=np.float32)  
    dx = np.diff(xs)  
    signs = np.sign(dx)  
    signs[np.abs(dx) < 1e-4] = 0  
    reversals, last = 0, 0  
    for s in signs:  
        if s == 0:  
            continue  
        if last == 0:  
            last = s  
        elif s != last:  
            reversals += 1  
            last = s  
    amplitude = float(xs.max() - xs.min())  
    return reversals, amplitude  

class JumpDetector:  
    """  
    简单双状态机 jump 检测：  
    - 使用踝部（左右取平均或可见侧）相对肩中心的归一化 y 坐标  
    - 速度低通，避免噪声  
    - rising -> falling 转换或回落确认时触发  
    """  
    def __init__(self):  
        self.state = "idle"  
        self.last_y = None  
        self.last_t = None  
        self.v_lp = 0.0  # 低通速度  
        self.y_min = None  

    def reset(self):  
        self.state = "idle"  
        self.last_y = None  
        self.last_t = None  
        self.v_lp = 0.0  
        self.y_min = None  

    def update(self, t_now, ankle_y_norm, cooldown_ok: bool):  
        if self.last_y is None or self.last_t is None:  
            self.last_y, self.last_t = ankle_y_norm, t_now  
            return None  

        dt = max(1e-3, t_now - self.last_t)  
        v = (ankle_y_norm - self.last_y) / dt  # y 向下为正，向上为负  
        # 速度低通  
        alpha = 1.0 - np.exp(-dt / max(1e-3, JUMP_SMOOTH_T))  
        self.v_lp = (1 - alpha) * self.v_lp + alpha * v  

        if self.state == "idle":  
            if self.v_lp < JUMP_VEL_UP_THRESH and cooldown_ok:  
                self.state = "rising"  
                self.y_min = ankle_y_norm  
        elif self.state == "rising":  
            self.y_min = min(self.y_min, ankle_y_norm)  
            if self.v_lp > 0.0:  
                self.state = "falling"  
        elif self.state == "falling":  
            ascent = (self.last_y - self.y_min)  
            if self.v_lp > JUMP_FALL_CONFIRM_VEL and ascent >= JUMP_MIN_ASCENT:  
                self.state = "idle"  
                self.y_min = None  
                self.last_y, self.last_t = ankle_y_norm, t_now  
                return {"ascent": float(ascent), "v_lp": float(self.v_lp)}  
            # 超时保护：若长时间没满足条件，回到 idle  
            if (t_now - self.last_t) > 1.0:  
                self.state = "idle"  
                self.y_min = None  

        self.last_y, self.last_t = ankle_y_norm, t_now  
        return None  

# ============================ 摄像头选择工具 ============================  
def find_working_camera(max_index=10, prefer=None, backend_flag=0):  
    """  
    探测可用摄像头索引。  
    prefer: 优先尝试的索引列表（如 [1, 0, 2]）  
    backend_flag: 可指定 OpenCV 后端（如 cv2.CAP_DSHOW / CAP_MSMF / CAP_V4L2 / CAP_AVFOUNDATION）  
    """  
    order = prefer if prefer is not None else range(max_index)  
    for idx in order:  
        cap = cv2.VideoCapture(idx, backend_flag) if backend_flag else cv2.VideoCapture(idx)  
        ok, _ = cap.read()  
        cap.release()  
        if ok:  
            return idx  
    return None 
    

def main():
    # 启动 WS 服务器
    ws_server = WSServer(HOST, PORT, source_id=SOURCE_ID)
    ws_server.start()

    # 选择外接摄像头：优先尝试 1，失败回退自动扫描；按平台设定后端
    import platform
    system = platform.system()
    if system == "Windows":
        backend = cv2.CAP_DSHOW  # 或尝试 cv2.CAP_MSMF
    elif system == "Linux":
        backend = cv2.CAP_V4L2
    elif system == "Darwin":
        backend = cv2.CAP_AVFOUNDATION
    else:
        backend = 0

    preferred_index = 1  # 外接摄像头通常为 1
    cap = cv2.VideoCapture(preferred_index, backend) if backend else cv2.VideoCapture(preferred_index)
    ok, _ = (cap.read() if cap.isOpened() else (False, None))
    if not ok:
        if cap:
            cap.release()
        print(f"[CAM] Preferred camera {preferred_index} open failed, scanning...")
        idx = find_working_camera(max_index=10, prefer=[1] + list(range(0, 10)), backend_flag=backend)
        if idx is None:
            print("Failed to open camera.")
            return
        print(f"[CAM] Using camera index: {idx}")
        cap = cv2.VideoCapture(idx, backend) if backend else cv2.VideoCapture(idx)

    # 可选：设置期望分辨率/帧率
    cap.set(cv2.CAP_PROP_FRAME_WIDTH,  1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    cap.set(cv2.CAP_PROP_FPS, 30)

    print(f"[INFO] WS server at ws://{HOST}:{PORT} | ESC to quit | MIRROR={MIRROR_PREVIEW}")

    pose = mp_pose.Pose(
        static_image_mode=False,
        model_complexity=1,  # 如需更稳可改 2
        smooth_landmarks=True,
        enable_segmentation=False,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5
    )

    prev_pts = None
    frame_count = 0

    # 持续帧与冷却
    hold = {"squat": 0}
    last_trigger = {"squat": 0.0, "jump": 0.0, "wave_left": 0.0, "wave_right": 0.0}

    # 挥手历史
    wave_hist = {"left": deque(), "right": deque()}

    # 跳跃检测器
    jump_det = JumpDetector()

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frame_count += 1

            if MIRROR_PREVIEW:
                frame = cv2.flip(frame, 1)

            h, w = frame.shape[:2]
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            res = pose.process(rgb)

            info = f"WS ws://{HOST}:{PORT} | MIRROR={MIRROR_PREVIEW}"
            cv2.putText(frame, info, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 2)

            if res.pose_landmarks:
                lm = res.pose_landmarks.landmark

                # 像素坐标 + EMA 平滑
                pts = compute_keypoints_xy(lm, w, h)
                pts = ema(prev_pts, pts, EMA_ALPHA)
                prev_pts = pts

                # 可见度数组
                vis = np.array([getattr(l, "visibility", 0.0) for l in lm], dtype=np.float32)
                def vis_ok(i, thr=VIS_THRESH): return vis[i] >= thr
                def P(i): return pts[i]

                # 肩中心与肩宽
                sh_center, shw = shoulder_center_and_width(pts)
                if shw < 5.0:
                    cv2.imshow("Pose (MediaPipe)", frame)
                    if cv2.waitKey(1) & 0xFF == 27:
                        break
                    continue

                # 归一化函数（肩中心为原点，除以肩宽）
                def norm(p):
                    return (p - sh_center) / shw

                # 取用关键点快捷名
                L = LMS
                l_sh = P(L.LEFT_SHOULDER); r_sh = P(L.RIGHT_SHOULDER)
                l_el = P(L.LEFT_ELBOW);    r_el = P(L.RIGHT_ELBOW)
                l_wr = P(L.LEFT_WRIST);    r_wr = P(L.RIGHT_WRIST)
                l_hp = P(L.LEFT_HIP);      r_hp = P(L.RIGHT_HIP)
                l_kn = P(L.LEFT_KNEE);     r_kn = P(L.RIGHT_KNEE)
                l_an = P(L.LEFT_ANKLE);    r_an = P(L.RIGHT_ANKLE)

                l_wr_n = norm(l_wr); r_wr_n = norm(r_wr)
                l_el_n = norm(l_el); r_el_n = norm(r_el)
                l_sh_n = norm(l_sh); r_sh_n = norm(r_sh)
                l_hp_n = norm(l_hp); r_hp_n = norm(r_hp)
                l_kn_n = norm(l_kn); r_kn_n = norm(r_kn)
                l_an_n = norm(l_an); r_an_n = norm(r_an)

                # 可见性检查
                left_knee_ok  = vis_ok(L.LEFT_KNEE)
                right_knee_ok = vis_ok(L.RIGHT_KNEE)
                left_hip_ok   = vis_ok(L.LEFT_HIP)
                right_hip_ok  = vis_ok(L.RIGHT_HIP)
                left_ankle_ok = vis_ok(L.LEFT_ANKLE)
                right_ankle_ok= vis_ok(L.RIGHT_ANKLE)

                knees_ok = left_knee_ok and right_knee_ok
                hips_ok  = left_hip_ok and right_hip_ok
                any_ankle_ok = left_ankle_ok or right_ankle_ok

                # ====== 可视化：显示更多关键点并区分可见度（含膝盖）======
                draw_ids = [
                    L.LEFT_SHOULDER, L.RIGHT_SHOULDER,
                    L.LEFT_ELBOW, L.RIGHT_ELBOW,
                    L.LEFT_WRIST, L.RIGHT_WRIST,
                    L.LEFT_HIP, L.RIGHT_HIP,
                    L.LEFT_KNEE, L.RIGHT_KNEE,
                    L.LEFT_ANKLE, L.RIGHT_ANKLE,
                ]
                for idx in draw_ids:
                    p = pts[idx]
                    vis_i = getattr(lm[idx], "visibility", 0.0)
                    color = (0, 255, 0) if vis_i >= VIS_THRESH else (128, 128, 128)
                    cv2.circle(frame, (int(p[0]), int(p[1])), 4, color, -1)

                # 画骨架线
                bones = [
                    (L.LEFT_SHOULDER, L.LEFT_ELBOW), (L.LEFT_ELBOW, L.LEFT_WRIST),
                    (L.RIGHT_SHOULDER, L.RIGHT_ELBOW), (L.RIGHT_ELBOW, L.RIGHT_WRIST),
                    (L.LEFT_HIP, L.LEFT_KNEE), (L.LEFT_KNEE, L.LEFT_ANKLE),
                    (L.RIGHT_HIP, L.RIGHT_KNEE), (L.RIGHT_KNEE, L.RIGHT_ANKLE),
                    (L.LEFT_SHOULDER, L.RIGHT_SHOULDER),
                    (L.LEFT_HIP, L.RIGHT_HIP),
                    (L.LEFT_SHOULDER, L.LEFT_HIP), (L.RIGHT_SHOULDER, L.RIGHT_HIP),
                ]
                for a, b in bones:
                    pa, pb = pts[a], pts[b]
                    cv2.line(frame, (int(pa[0]), int(pa[1])), (int(pb[0]), int(pb[1])), (200, 200, 200), 1)

                # ================= 下蹲（融合多特征 + 降级策略） =================
                squat_now = False
                squat_details = None

                if knees_ok and hips_ok:
                    # 髋、膝 y（归一化）
                    hip_y   = 0.5 * (l_hp_n[1] + r_hp_n[1])
                    knee_y  = 0.5 * (l_kn_n[1] + r_kn_n[1])

                    # 膝角：各侧需要 hip/knee/ankle 都可见
                    knee_angles = []
                    if left_ankle_ok:
                        knee_angles.append(angle_deg(pts[L.LEFT_HIP], pts[L.LEFT_KNEE], pts[L.LEFT_ANKLE]))
                    if right_ankle_ok:
                        knee_angles.append(angle_deg(pts[L.RIGHT_HIP], pts[L.RIGHT_KNEE], pts[L.RIGHT_ANKLE]))
                    knee_angle_ok = (len(knee_angles) >= 1 and all(ka < SQUAT_KNEE_ANGLE_MAX for ka in knee_angles))

                    knee_proximity_ok = (hip_y > knee_y - SQUAT_KNEE_PROX_DELTA)

                    hip_ankle_ok = False
                    ankle_y = None
                    hip_ankle_gap = None
                    if any_ankle_ok:
                        # 用可见踝的均值作为踝 y
                        y_list = []
                        if left_ankle_ok:  y_list.append(l_an_n[1])
                        if right_ankle_ok: y_list.append(r_an_n[1])
                        ankle_y = float(np.mean(y_list))
                        hip_ankle_gap = hip_y - ankle_y
                        hip_ankle_ok = (hip_ankle_gap < SQUAT_HIP_ANKLE_GAP_MAX)

                    # 辅助：肩 y
                    sh_y = 0.5 * (l_sh_n[1] + r_sh_n[1])
                    hip_drop_vs_shoulder_ok = (hip_y > sh_y + SQUAT_HIP_DROP_VS_SHOULDER)

                    votes = [knee_proximity_ok, knee_angle_ok]
                    if any_ankle_ok:
                        votes.append(hip_ankle_ok)
                    strong_votes = sum(votes)

                    if any_ankle_ok:
                        squat_now = (strong_votes >= 2) or (knee_angle_ok and hip_drop_vs_shoulder_ok)
                    else:
                        # 无踝时仅使用两票（接近+角度）
                        squat_now = (strong_votes >= 2)

                    squat_details = {
                        "hip_y": float(hip_y),
                        "knee_y": float(knee_y),
                        "ankle_y": None if ankle_y is None else float(ankle_y),
                        "knee_angles": list(map(float, knee_angles)),
                        "gap_hip_ankle": None if hip_ankle_gap is None else float(hip_ankle_gap),
                        "votes": int(strong_votes),
                        "knees_ok": bool(knees_ok),
                        "hips_ok": bool(hips_ok),
                        "any_ankle_ok": bool(any_ankle_ok)
                    }

                # 触发（保持你的持有与冷却逻辑）
                t_now = time.time()
                hold["squat"] = hold["squat"] + 1 if squat_now else 0
                if hold["squat"] >= SQUAT_HOLD_FRAMES and (t_now - last_trigger["squat"]) > SQUAT_COOLDOWN_S:
                    ws_server.send_event("squat", details=squat_details)
                    last_trigger["squat"] = t_now
                    hold["squat"] = 0

                # ================= 跳跃（需至少一侧踝可见） =================
                jump_evt = None
                if any_ankle_ok:
                    an_y_list = []
                    if left_ankle_ok:  an_y_list.append(l_an_n[1])
                    if right_ankle_ok: an_y_list.append(r_an_n[1])
                    an_y = float(np.mean(an_y_list))

                    jump_cool_ok = (t_now - last_trigger["jump"]) > JUMP_COOLDOWN_S
                    jump_evt = jump_det.update(t_now, an_y, jump_cool_ok)

                if jump_evt is not None:
                    ws_server.send_event("jump", details=jump_evt)
                    last_trigger["jump"] = t_now

                # ================= 挥手 =================
                rev_l, amp_l = update_wave(wave_hist["left"], t_now, float(l_wr_n[0]))
                rev_r, amp_r = update_wave(wave_hist["right"], t_now, float(r_wr_n[0]))
                if (t_now - last_trigger["wave_left"]) > WAVE_COOLDOWN_S:
                    if rev_l >= WAVE_MIN_REVERSALS and amp_l >= WAVE_MIN_AMPLITUDE_SHOULDER:
                        ws_server.send_event("wave_left", details={"reversals": int(rev_l), "amplitude": float(amp_l)})
                        last_trigger["wave_left"] = t_now
                        wave_hist["left"].clear()
                if (t_now - last_trigger["wave_right"]) > WAVE_COOLDOWN_S:
                    if rev_r >= WAVE_MIN_REVERSALS and amp_r >= WAVE_MIN_AMPLITUDE_SHOULDER:
                        ws_server.send_event("wave_right", details={"reversals": int(rev_r), "amplitude": float(amp_r)})
                        last_trigger["wave_right"] = t_now
                        wave_hist["right"].clear()

                # ================= 可视化中心点（保留） =================
                cx, cy = int(sh_center[0]), int(sh_center[1])
                cv2.circle(frame, (cx, cy), 5, (0, 255, 0), -1)

            cv2.imshow("Pose (MediaPipe)", frame)
            if cv2.waitKey(1) & 0xFF == 27:
                break

    finally:
        cap.release()
        cv2.destroyAllWindows()
        

if __name__ == "__main__":
    main()